#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
百度网盘上传操作模块
提供文件上传和秒传功能
"""

import os
import json
import time
import hashlib
import logging
import requests
from typing import Dict, List, Optional, Any, Union, BinaryIO


class UploadOperations:
    """百度网盘上传操作类"""
    
    def __init__(self, client):
        """初始化上传操作
        
        Args:
            client: 百度网盘客户端实例
        """
        self.client = client
        self.logger = client.logger
        self.session = client.session
        
        # 百度网盘API端点
        self.api_base = "https://pan.baidu.com/rest/2.0/xpan"
        self.upload_base = "https://d.pcs.baidu.com/rest/2.0/pcs"
        
        # 分片大小（4MB）
        self.chunk_size = 4 * 1024 * 1024
    
    def check_fast_upload(self, filename: str, file_size: int, file_md5: str, 
                         parent_path: str = "/") -> bool:
        """检查文件是否支持秒传
        
        Args:
            filename: 文件名
            file_size: 文件大小（字节）
            file_md5: 文件MD5值
            parent_path: 父文件夹路径
        
        Returns:
            bool: 是否支持秒传
        """
        if not self.client.check_login_status():
            return False
        
        file_path = os.path.join(parent_path, filename).replace('\\', '/')
        
        url = f"{self.api_base}/file"
        params = {
            'method': 'rapidupload',
            'access_token': self.client.access_token
        }
        data = {
            'path': file_path,
            'content-md5': file_md5,
            'slice-md5': file_md5,  # 简化处理，实际应该是第一个分片的MD5
            'content-length': file_size
        }
        
        try:
            response = self.session.post(url, params=params, data=data)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    self.logger.info(f"文件支持秒传: {filename}")
                    return True
                elif result.get('errno') == 31079:
                    # 文件不存在，不支持秒传
                    return False
                else:
                    self.logger.warning(f"检查秒传失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"检查秒传失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"检查秒传请求失败: {e}")
        
        return False
    
    def fast_upload(self, filename: str, file_size: int, file_md5: str, 
                   parent_path: str = "/") -> Optional[Dict[str, Any]]:
        """执行秒传
        
        Args:
            filename: 文件名
            file_size: 文件大小（字节）
            file_md5: 文件MD5值
            parent_path: 父文件夹路径
        
        Returns:
            Dict[str, Any]: 秒传成功返回文件信息，失败返回None
        """
        if not self.client.check_login_status():
            return None
        
        file_path = os.path.join(parent_path, filename).replace('\\', '/')
        
        url = f"{self.api_base}/file"
        params = {
            'method': 'rapidupload',
            'access_token': self.client.access_token
        }
        data = {
            'path': file_path,
            'content-md5': file_md5,
            'slice-md5': file_md5,  # 简化处理
            'content-length': file_size
        }
        
        try:
            response = self.session.post(url, params=params, data=data)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    self.logger.info(f"文件秒传成功: {filename}")
                    return {
                        'id': str(result.get('fs_id')),
                        'name': filename,
                        'path': file_path,
                        'size': file_size,
                        'type': 'file',
                        'md5': file_md5,
                        'create_time': int(time.time()),
                        'parent_path': parent_path
                    }
                else:
                    self.logger.error(f"文件秒传失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"文件秒传失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"文件秒传请求失败: {e}")
        
        return None
    
    def upload_file(self, file_data: Union[str, bytes, BinaryIO], filename: str, 
                   parent_path: str = "/", overwrite: bool = True) -> Optional[Dict[str, Any]]:
        """上传文件
        
        Args:
            file_data: 文件数据（文件路径、字节数据或文件对象）
            filename: 文件名
            parent_path: 父文件夹路径
            overwrite: 是否覆盖同名文件
        
        Returns:
            Dict[str, Any]: 上传成功返回文件信息，失败返回None
        """
        if not self.client.check_login_status():
            return None
        
        # 获取文件数据和大小
        if isinstance(file_data, str):
            # 文件路径
            file_size = os.path.getsize(file_data)
            with open(file_data, 'rb') as f:
                file_content = f.read()
        elif isinstance(file_data, bytes):
            # 字节数据
            file_content = file_data
            file_size = len(file_content)
        else:
            # 文件对象
            file_content = file_data.read()
            file_size = len(file_content)
        
        # 计算MD5
        file_md5 = hashlib.md5(file_content).hexdigest()
        
        # 尝试秒传
        if self.check_fast_upload(filename, file_size, file_md5, parent_path):
            result = self.fast_upload(filename, file_size, file_md5, parent_path)
            if result:
                return result
        
        # 普通上传
        return self._normal_upload(file_content, filename, parent_path, overwrite)
    
    def _normal_upload(self, file_content: bytes, filename: str, 
                      parent_path: str = "/", overwrite: bool = True) -> Optional[Dict[str, Any]]:
        """普通文件上传
        
        Args:
            file_content: 文件内容
            filename: 文件名
            parent_path: 父文件夹路径
            overwrite: 是否覆盖同名文件
        
        Returns:
            Dict[str, Any]: 上传成功返回文件信息，失败返回None
        """
        file_size = len(file_content)
        file_path = os.path.join(parent_path, filename).replace('\\', '/')
        
        # 小文件直接上传
        if file_size <= self.chunk_size:
            return self._upload_small_file(file_content, file_path, overwrite)
        else:
            # 大文件分片上传
            return self._upload_large_file(file_content, file_path, overwrite)
    
    def _upload_small_file(self, file_content: bytes, file_path: str, 
                          overwrite: bool = True) -> Optional[Dict[str, Any]]:
        """上传小文件
        
        Args:
            file_content: 文件内容
            file_path: 文件路径
            overwrite: 是否覆盖同名文件
        
        Returns:
            Dict[str, Any]: 上传成功返回文件信息，失败返回None
        """
        # 预上传
        upload_info = self._pre_upload(file_path, len(file_content), overwrite)
        if not upload_info:
            return None
        
        upload_url = upload_info.get('uploadurl')
        if not upload_url:
            self.logger.error("获取上传URL失败")
            return None
        
        # 上传文件
        try:
            files = {
                'file': (os.path.basename(file_path), file_content, 'application/octet-stream')
            }
            
            response = self.session.post(upload_url, files=files)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    # 创建文件
                    return self._create_file(file_path, len(file_content), 
                                           result.get('md5'), upload_info.get('uploadid'))
                else:
                    self.logger.error(f"上传文件失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"上传文件失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"上传文件请求失败: {e}")
        
        return None
    
    def _upload_large_file(self, file_content: bytes, file_path: str, 
                          overwrite: bool = True) -> Optional[Dict[str, Any]]:
        """上传大文件（分片上传）
        
        Args:
            file_content: 文件内容
            file_path: 文件路径
            overwrite: 是否覆盖同名文件
        
        Returns:
            Dict[str, Any]: 上传成功返回文件信息，失败返回None
        """
        file_size = len(file_content)
        
        # 预上传
        upload_info = self._pre_upload(file_path, file_size, overwrite)
        if not upload_info:
            return None
        
        upload_url = upload_info.get('uploadurl')
        upload_id = upload_info.get('uploadid')
        
        if not upload_url or not upload_id:
            self.logger.error("获取上传信息失败")
            return None
        
        # 分片上传
        chunk_list = []
        chunk_num = 0
        
        for i in range(0, file_size, self.chunk_size):
            chunk_data = file_content[i:i + self.chunk_size]
            chunk_md5 = hashlib.md5(chunk_data).hexdigest()
            
            # 上传分片
            if self._upload_chunk(upload_url, chunk_data, chunk_num):
                chunk_list.append(chunk_md5)
                chunk_num += 1
            else:
                self.logger.error(f"上传分片 {chunk_num} 失败")
                return None
        
        # 创建文件
        file_md5 = hashlib.md5(file_content).hexdigest()
        return self._create_file(file_path, file_size, file_md5, upload_id, chunk_list)
    
    def _pre_upload(self, file_path: str, file_size: int, 
                   overwrite: bool = True) -> Optional[Dict[str, Any]]:
        """预上传
        
        Args:
            file_path: 文件路径
            file_size: 文件大小
            overwrite: 是否覆盖同名文件
        
        Returns:
            Dict[str, Any]: 预上传信息
        """
        url = f"{self.api_base}/file"
        params = {
            'method': 'precreate',
            'access_token': self.client.access_token
        }
        data = {
            'path': file_path,
            'size': file_size,
            'isdir': 0,
            'rtype': 1 if overwrite else 0
        }
        
        try:
            response = self.session.post(url, params=params, data=data)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    return result
                else:
                    self.logger.error(f"预上传失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"预上传失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"预上传请求失败: {e}")
        
        return None
    
    def _upload_chunk(self, upload_url: str, chunk_data: bytes, chunk_num: int) -> bool:
        """上传分片
        
        Args:
            upload_url: 上传URL
            chunk_data: 分片数据
            chunk_num: 分片序号
        
        Returns:
            bool: 上传是否成功
        """
        try:
            files = {
                'file': (f'chunk_{chunk_num}', chunk_data, 'application/octet-stream')
            }
            data = {
                'partseq': chunk_num
            }
            
            response = self.session.post(upload_url, files=files, data=data)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    return True
                else:
                    self.logger.error(f"上传分片失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"上传分片失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"上传分片请求失败: {e}")
        
        return False
    
    def _create_file(self, file_path: str, file_size: int, file_md5: str, 
                    upload_id: str, chunk_list: List[str] = None) -> Optional[Dict[str, Any]]:
        """创建文件
        
        Args:
            file_path: 文件路径
            file_size: 文件大小
            file_md5: 文件MD5
            upload_id: 上传ID
            chunk_list: 分片MD5列表
        
        Returns:
            Dict[str, Any]: 创建成功返回文件信息，失败返回None
        """
        url = f"{self.api_base}/file"
        params = {
            'method': 'create',
            'access_token': self.client.access_token
        }
        data = {
            'path': file_path,
            'size': file_size,
            'uploadid': upload_id,
            'isdir': 0,
            'rtype': 1
        }
        
        if chunk_list:
            data['block_list'] = json.dumps(chunk_list)
        
        try:
            response = self.session.post(url, params=params, data=data)
            if response.status_code == 200:
                result = response.json()
                if result.get('errno') == 0:
                    self.logger.info(f"文件上传成功: {file_path}")
                    return {
                        'id': str(result.get('fs_id')),
                        'name': os.path.basename(file_path),
                        'path': file_path,
                        'size': file_size,
                        'type': 'file',
                        'md5': file_md5,
                        'create_time': int(time.time()),
                        'parent_path': os.path.dirname(file_path)
                    }
                else:
                    self.logger.error(f"创建文件失败: {result.get('errmsg')}")
            else:
                self.logger.error(f"创建文件失败，状态码: {response.status_code}")
        except Exception as e:
            self.logger.error(f"创建文件请求失败: {e}")
        
        return None
    
    def calculate_slice_md5(self, file_data: Union[str, bytes, BinaryIO], 
                           slice_size: int = None) -> str:
        """计算文件分片MD5值
        
        Args:
            file_data: 文件数据（文件路径、字节数据或文件对象）
            slice_size: 分片大小，None表示使用默认大小
        
        Returns:
            str: 第一个分片的MD5值
        """
        if slice_size is None:
            slice_size = self.chunk_size
        
        if isinstance(file_data, str):
            # 文件路径
            with open(file_data, 'rb') as f:
                chunk = f.read(slice_size)
        elif isinstance(file_data, bytes):
            # 字节数据
            chunk = file_data[:slice_size]
        else:
            # 文件对象
            chunk = file_data.read(slice_size)
        
        return hashlib.md5(chunk).hexdigest()